# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

import enum
import logging
import math
import typing

import requests
import tornado.gen
import logging
import coin.flow.subscriber as flow_subscriber
import coin.flow.topic_map as topic_map

import coin.proto.coin_order_enums_pb2 as coin_order_pb2
from coin.base.config import Config
from coin.exchange.base.order_fill_util import FillManager
from coin.exchange.base.rest.task_scheduler import Scheduler, wrap_async_task
from coin.exchange.util.rest_rate_limit import RestRateLimiter
from coin.proto.coin_order_gateway_pb2 import OrderGatewayConfig
from coin.proto.coin_query_pb2 import AccountBalance
from coin.exchange.uniswap_v3.kr_rest.currency import UniswapCurrency
from coin.exchange.uniswap_v3.kr_rest.product import UniswapProduct
from coin.exchange.uniswap_v3.kr_rest.private_client_v2 import (
    UniswapPrivateClient,
    UniswapPrivateParser,
)
from coin.exchange.uniswap_v3.kr_rest.native_private_client import (
    from_uniswap_order_side,
    from_uniswap_order_type,
    to_uniswap_order_side,
    to_uniswap_order_type,
    UniswapOrderSide,
    UniswapOrderType,
    UniswapOrderStatus,
    TxReceipt,
)
from coin.exchange.base.rest_order_gateway_base import (
    RestOrderGateway,
    exchange_message_handler,
    create_order_internal,
    OrderEvent,
    Order,
    OrderType,
    BaseBalanceOrderGatewayInfo,
    PrivateExchangeMessage,
    ProductFillElement,
    to_proto_order_side,
    coin_order_pb2,
    pretty_json,
)
from coin.proto.coin_query_pb2 import (
    AccountBalance,
    CurrencyBalance,
)

import threading
import time


def create_order_from_order_info(order_msg, timestamp, order_id=None):
    price_orig = float(order_msg['price'])
    qty_orig = float(order_msg['quantity'])
    order_side = from_uniswap_order_side(UniswapOrderSide(order_msg['side']))
    order_type = from_uniswap_order_type(UniswapOrderType(order_msg['order_type']))
    exchange_order_id = str(order_msg['id'])
    product = UniswapProduct.FromStrNativeProduct(order_msg['currency_pair_code'])
    status = UniswapOrderStatus(order_msg['status'])
    if order_msg['order_fee'] is None:
        fee = 0
    else:
        fee = float(order_msg['order_fee'])

    fill = FillManager(
        qty_orig=qty_orig,
        price_orig=price_orig,
        qty_accum_fill=0.0,
        price_avg_fill=None,
        qty_last_fill=None,
        price_last_fill=None,
    )

    executions = []
    fill_id_set = set()  # Sometimes Quoinex gives duplicated fill msgs.
    for execution_info in order_msg.get('executions', []):
        fill_id = '%s_%s' % (product.symbol, execution_info['id'])
        if fill_id in fill_id_set:
            # Ignore
            continue
        else:
            fill_id_set.add(fill_id)

        price_last_fill = float(execution_info['price'])
        qty_last_fill = float(execution_info['quantity'])
        fill.update_by_last_fill(price_last_fill, qty_last_fill)
        taker_side = execution_info['taker_side']
        my_side = execution_info['my_side']

        if my_side == taker_side:
            fill_type = coin_order_pb2.TAKER_FILL_TYPE
        else:
            fill_type = coin_order_pb2.MAKER_FILL_TYPE

        executions.append({
            'price': price_last_fill,
            'qty': qty_last_fill,
            'fill_id': fill_id,
            'fill_type': fill_type,
        })

    qty_accum_fill = float(order_msg['filled_quantity'])
    if not math.isclose(fill.value.qty_accum_fill, qty_accum_fill):
        logging.error("order info's filled_qty is not consistent! %s", pretty_json(order_msg))

    order = Order(order_id=order_id,
                  product=product,
                  price=price_orig,
                  qty=qty_orig,
                  order_side=order_side,
                  order_type=order_type,
                  tag={'executions': executions},
                  accepted=True,
                  cancel_sent=False,
                  internal=create_order_internal(
                      timestamp=timestamp,
                      exchange_order_id=exchange_order_id,
                      fill=fill,
                      status=status,
                      fee=fee,
                  ))
    return order


class UniswapOrderGatewayInfo(BaseBalanceOrderGatewayInfo):
    CurrencyType = UniswapCurrency
    ProductType = UniswapProduct

    def __init__(
        self,
        currencies,  # : List[CurrencyType],
        products  # : List[ProductType]):
    ):
        super().__init__(currencies=currencies, products=products)
        self._nft_token_info = {product.native_symbol: None for product in self.products}

    def update_balance(self, update_msg):
        account_balance = update_msg['msg']
        currencies_to_follow = [currency.currency for currency in self.currencies]
        currencies_to_follow.append('WETH')
        for balance in account_balance.each_balance:
            currency = UniswapCurrency.FromStrCurrency(balance.currency)
            if balance.currency in currencies_to_follow:
                self.set_balance(currency, balance.total, balance)
        eth_internal_balance = self._balances_internal.get('ETH', None)
        weth_internal_balance = self._balances_internal.get('WETH', None)
        eth_total = 0.0
        eth_available = 0.0
        eth_hold = 0.0
        if eth_internal_balance:
            eth_total += eth_internal_balance.total
            eth_available += eth_internal_balance.available
            eth_hold += eth_internal_balance.hold
        if weth_internal_balance:
            eth_total += weth_internal_balance.total
            eth_available += weth_internal_balance.available
            eth_hold += weth_internal_balance.hold
        self.set_balance(
            UniswapCurrency.FromStrCurrency('ETH'),
            eth_total,
            CurrencyBalance(
                currency='ETH',
                currency_native='ETH',
                available=eth_available,
                hold=eth_hold,
                total=eth_total,
            )
        )

    def balance_as_proto(self):
        return AccountBalance(exchange='Uniswap',
                              market_type='Spot',
                              each_balance=self.get_notNone_balances_internal())

    def get_balance_wrapped_only(self, currency):
        wrapped = f'W{currency}'
        if wrapped in self._balances.keys():
            return self.get_balance(wrapped)
        else:
            return self.get_balance(currency)

    def update_nft_token_info(self, update_msg):
        for symbol, nft_token_info_response in update_msg.items():
            self._nft_token_info[symbol] = nft_token_info_response.get('tokenid_to_nft_info', {})

    def get_nft_token_info(self, symbol):
        return self._nft_token_info[symbol]


class UniswapOrderGateway(RestOrderGateway):
    _exchange = 'Uniswap'
    _market_type = 'Spot'
    _api_version = 'v3'

    OGInfoType = UniswapOrderGatewayInfo

    class _MsgType(enum.IntEnum):
        REST_ACCOUNT_BALANCE = 1
        REST_RECENT_ORDERS = 2
        REST_NFT_TOKEN_INFO = 3

    def __init__(self,
                 currencies: typing.List[UniswapCurrency],
                 products: typing.List[UniswapProduct],
                 config: Config,
                 og_config=None,
                 logger=None,
                 blockchain_name="ethereum",
                 api_version="v3"):
        self._api_version = api_version
        super().__init__(currencies, products, config, og_config, logger)

        self._blockchain_name = blockchain_name
        self._private_client = None
        self._scheduler = None
        self._rest_rate_limiter = RestRateLimiter(delay=300)

        self._ioloop = None
        self._flow_sub = None
        self._publisher = None
        self._fire_lock_map = {product.symbol: threading.Lock() for product in products}

        self._og_config = OrderGatewayConfig(
            account_query_period=10,
            orders_snapshot_query_period=10,
            use_fill_checker=False,
        )
        if og_config is not None:
            self._og_config.MergeFrom(og_config)

    def get_balance_wrapped_only(self, currency) -> float:
        return self._og_info.get_balance_wrapped_only(currency)

    def submittable(self):
        return self.is_ready() and self._rest_rate_limiter.is_available()

    def start(self, queue, ioloop, flow_sub=None):
        api_host = self._og_config.rest_root_url \
            if self._og_config.HasField('rest_root_url') else None
        mea = f'{self._market_type}.{self._exchange}.{self._api_version}'
        self._private_client = UniswapPrivateClient(
            wallet_address=self._config.owner,
            og_info=self._og_info,
            blockchain_name=self._blockchain_name,
            mea=mea,
        )
        topic_id = topic_map.search_topic_id('uniswap_rest_private_1')[0][1]
        self._ioloop = ioloop
        self._publisher = queue.get_writer(topic_id=topic_id)
        self._scheduler = Scheduler(ioloop=ioloop, logger=self._logger)
        self._install_periodic_query()

        # Subscribe to flow for feed.
        self._flow_sub = flow_sub or flow_subscriber.from_queue(queue)
        self._flow_sub.subscribe('uniswap_rest_private_1', self.on_private_exchange_msg)

    @tornado.gen.coroutine
    def _install_periodic_query(self):
        assert self._scheduler is not None, 'self._scheduler is not set!'
        assert self._publisher is not None, 'self._publisher is not set!'

        # query balance
        query_balance_task = wrap_async_task(
            self._private_client.async_query_account_balance,
            msg_type=self._MsgType.REST_ACCOUNT_BALANCE,
            publisher=self._publisher,
        )
        self._scheduler.add_periodic_task(
            period=self._og_config.account_query_period,
            job_func=query_balance_task,
        )

        # query nft token info
        query_nft_token_info = wrap_async_task(
            self._private_client.async_query_nft_token_info,
            msg_type=self._MsgType.REST_NFT_TOKEN_INFO,
            publisher=self._publisher,
        )
        self._scheduler.add_periodic_task(
            period=self._og_config.account_query_period,
            job_func=query_nft_token_info,
        )
        yield tornado.gen.sleep(self._og_config.burst_query_avoid_time)
        self.start_scheduler()

    @tornado.gen.coroutine
    def _async_submit_order(self, order, post_only=False):
        if post_only:
            self._logger.warning('post only not supported.')
        assert order.order_id is not None, order
        assert order.order_type == OrderType.LIMIT, order
        product = order.product
        assert product is not None, order

        formatter = self._order_formatter[product.native_symbol]
        params = {
            'product': product,
            'price': formatter.format_price(order.price),
            'qty': formatter.format_qty(order.qty),
            'side': to_uniswap_order_side(order.order_side) if order.tag is None else order.tag,
        }
        return self._private_client.query_submit(**params)

    def _async_liquidity_pool_submit_order(self, order, **kwargs):
        assert order.order_id is not None, order
        assert order.order_type == OrderType.LIMIT, order
        product = order.product
        assert product is not None, order
        assert product.native_symbol in self._fire_lock_map.keys()

        lock = self._fire_lock_map[product.native_symbol]
        if lock.acquire(False):
            formatter = self._order_formatter[product.native_symbol]
            params = {
                'product': product,
                'price': formatter.format_price(order.price),
                'qty': formatter.format_qty(order.qty),
                'side': order.order_side,
                'execution_lock': lock,
            }
            return self._private_client.query_liquidity_pool_submit(**{**params, **kwargs})
        else:
            logging.info("order rejected")
            return None

    def _handle_submit_response(self, order, response: TxReceipt):
        logging.info(f'submit_response: order = {order} | response = {response}')

        order.accepted = True
        order.internal.fully_filled = True
        if response['status'] == 'success' or response['status'] == True:
            order.internal.status = UniswapOrderStatus.FILLED
            order.internal.exchange_order_id = response['tx_hash']
            order.tag = {
                'executions': [
                    {'fill_id': str(time.time())}
                ]
            }
            if response.get('fill_price') is not None:
                order.price = response['fill_price']
                self._update_order(order, order.internal.timestamp)
        else:
            order.internal.status = UniswapOrderStatus.CANCELED
        self._order_manager.remove_order(order)

    @exchange_message_handler(_MsgType.REST_ACCOUNT_BALANCE)
    def _handle_account_balance(self, exchange_message: PrivateExchangeMessage):
        response = exchange_message.data
        update_msg = response.json()
        self._og_info.update_balance(update_msg)
        self._og_logger.write_balance(self._og_info.balance_as_proto())
        if not self.is_ready():
            self._og_ready.set_balance_ready(True)
            self._og_ready.set_orders_ready(True)
            self._og_ready.check_ready()

    @exchange_message_handler(_MsgType.REST_NFT_TOKEN_INFO)
    def _handle_nft_token_info(self, exchange_message: PrivateExchangeMessage):
        response = exchange_message.data
        self._og_info.update_nft_token_info(response)
        if not self.is_ready():
            self._og_ready.set_nft_token_info_ready(True)
            self._og_ready.set_orders_ready(True)
            self._og_ready.check_ready()

    def get_nft_token_info(self, symbol):
        return self._og_info.get_nft_token_info(symbol)

    def _async_cancel_order(self, order):
        raise NotImplementedError()

    def _handle_cancel_response(self, order, response: requests.Response):
        raise NotImplementedError()

    def _update_order(self, order_info, timestamp, is_first_query=True):
        exchange_order_id = order_info.internal.exchange_order_id
        status = order_info.internal.status
        order = self._order_manager.get(exchange_order_id=exchange_order_id)

        if order is None:
            order = self._order_manager.get(order_id=order_info.order_id)

            # Handle foreign order.
            if order is None:
                self._order_manager.add_order(order_info)
                order = order_info
            
            if not order.accepted:
                self._order_manager.set_order(order, exchange_order_id=exchange_order_id)

        if status == UniswapOrderStatus.CANCELED:
            if not self._order_manager.is_dead_order(order):
                self._order_manager.set_order(order, cancel_confirmed=True)
        elif status == UniswapOrderStatus.FILLED:
            fill_proto = ProductFillElement(
                fill_id=order.order_id,
                price=order.price,
                qty=order.qty,
                order_id=order.order_id,
                side=to_proto_order_side(order.order_side),
                exchange_order_id=exchange_order_id,
                fill_type=coin_order_pb2.TAKER_FILL_TYPE,
                symbol=order.product.native_symbol)
            self.process_fill(fill_proto, timestamp, order)
            if not self._order_manager.is_dead_order(order):
                self._order_manager.set_order(order, fully_filled=True)
        elif status == UniswapOrderStatus.LIVE:
            if self._order_manager.is_dead_order(order):
                self._logger.info('Order is already marked as dead, %s', order.order_id)
            else:
                if order.internal.exchange_order_id != exchange_order_id:
                    # Should never happen
                    self._logger.error('Inconsistent exchange_order_id, %s\n%s', exchange_order_id, order)
        else:
            self._logger.error('Unknown order status: %s', status)

